/**
 * @author Gao Jian
 * @email gaoj@njupt.edu.cn
 * @create date 2021-06-03 10:38:11
 * @desc 膨胀与腐蚀
 */

#include "defs.h"

__global__ void kdilateparts(unsigned char *output, unsigned char *input,
                             unsigned char value, unsigned int nw,
                             unsigned int nh) {
  unsigned nz = nw * nh;
  unsigned increment = blockDim.x * gridDim.x;
  unsigned nits = nz / increment + 1;
  unsigned idx = blockDim.x * blockIdx.x + threadIdx.x;

  for (unsigned i = 0; i < nits; i++, idx += increment) {
    if (idx >= nz)
      break;

    if (input[idx] == value) {
      output[idx] = value;
      continue;
    }

    unsigned iw = idx % nw;
    unsigned ih = idx / nw;

    unsigned char bbond[8] = {iw > 0,
                              iw<nw - 1, ih> 0,
                              ih<nh - 1, iw> 0 && ih > 0,
                              iw > 0 && ih < nh - 1,
                              iw < nw - 1 && ih > 0,
                              iw < nw - 1 && ih < nh - 1};
    unsigned int ids[8] = {idx - 1,      idx + 1,      idx - nw,
                           idx + nw,     idx - 1 - nw, idx - 1 + nw,
                           idx + 1 - nw, idx + 1 + nw};

    for (unsigned char i = 0; i < 8; i++) {
      if (bbond[i] && input[ids[i]] == value) {
        output[idx] = value;
        break;
      }
    }
  }
}

int dilate_parts(unsigned char *data, unsigned char value, unsigned int nit,
                 unsigned int nw, unsigned int nh, const unsigned verbose = 0) {
  unsigned int nz = nw * nh;
  unsigned char *d_input = NULL;
  unsigned char *d_output = NULL;
  unsigned char *tmp = NULL;

  cudaMalloc((void **)&d_input, nz * sizeof(unsigned char));
  cudaMalloc((void **)&d_output, nz * sizeof(unsigned char));
  cudaMemcpy(d_input, data, nz * sizeof(unsigned char), cudaMemcpyHostToDevice);
  cudaMemcpy(d_output, d_input, nz * sizeof(unsigned char),
             cudaMemcpyDeviceToDevice);

  tmp = d_output;

  for (unsigned char i = 0; i < nit; i++) {
    kdilateparts<<<NUM_BLOCKS, NUM_THREADS>>>(d_output, d_input, value, nw, nh);
    cudaDeviceSynchronize();
    if (verbose)
      printf("+");
    if ((i + 1) % 50 == 0 && verbose)
      printf("  [%3d]\n  ", i + 1);
    tmp = d_output;
    d_output = d_input;
    d_input = tmp;
  }
  if (verbose)
    printf("  [%3d]\n", nit);

  cudaMemcpy(data, tmp, nz * sizeof(unsigned char), cudaMemcpyDeviceToHost);
  cudaFree(d_output);
  cudaFree(d_input);

  return 0;
}